import torch
import torch.nn.functional as F
import torch.nn as nn

class FAM(nn.Module):
    def __init__(self, k, k_out, need_x2=False, need_fuse=False):
        super(FAM, self).__init__()
        self.pools_sizes = [2,4,8]
        self.need_x2 = need_x2
        self.need_fuse = need_fuse
        pools, convs = [],[]
        for i in self.pools_sizes:
            pools.append(nn.AvgPool2d(kernel_size=i, stride=i))
            convs.append(nn.Conv2d(k, k, 3, 1, 1, bias=False))
        self.pools = nn.ModuleList(pools)
        self.convs = nn.ModuleList(convs)
        self.relu = nn.ReLU()
        self.conv1 = nn.Conv2d(k, k_out, 3, 1, 1, bias=False)
        self.conv2 = nn.Conv2d(k, k, 3, 1, 1, bias=False)
        self.conv_sum = nn.Conv2d(k, k_out, 3, 1, 1, bias=False)
        if self.need_fuse:
            self.conv_sum_c = nn.Conv2d(k_out, k_out, 3, 1, 1, bias=False)

    def forward(self, x, x2=None, x3=None):
        x_size = x.size()
        resl = x
        for i in range(len(self.pools_sizes)):
            y = self.convs[i](self.pools[i](x))
            resl = torch.add(resl, F.interpolate(y, x_size[2:], mode='bilinear', align_corners=True))
        resl = self.relu(resl)
        if self.need_x2:
            resl = self.conv2(resl)
            resl = F.interpolate(resl, x2.size()[2:], mode='bilinear', align_corners=True)
            resl = torch.add(resl, x2)
            resl = self.conv_sum(resl)
        else:
            resl = self.conv1(resl)
        
        if self.need_fuse:
            resl = self.conv_sum_c(torch.add(resl, x3))
        return resl
    
if __name__ == '__main__':
    a = torch.randn(32,64,44,44)
    test = FAM(64,64)
    c = test(a)
    print(a.shape)